#pragma once

#include <rclcpp/rclcpp.hpp>
#include "NvInferPlugin.h"
#include "cuda_runtime_api.h"

inline int get_size_by_dims(const nvinfer1::Dims& dims)
{
    auto logger = rclcpp::get_logger("NvLogger");

    int size = 1;
    for (int i = 0; i < dims.nbDims; i++) {
        size *= dims.d[i];
        RCLCPP_INFO(logger, "dims: %d, %d", i, dims.d[i]);
    }
    return size;
}

struct Binding {
    size_t size = 1;
    size_t dsize = 1;
    nvinfer1::Dims dims;
    std::string name;
};

inline int type_to_size(const nvinfer1::DataType& dataType)
{
    switch (dataType) {
    case nvinfer1::DataType::kFLOAT:
        return 4;
    case nvinfer1::DataType::kHALF:
        return 2;
    case nvinfer1::DataType::kINT32:
        return 4;
    case nvinfer1::DataType::kINT8:
        return 1;
    case nvinfer1::DataType::kBOOL:
        return 1;
    default:
        return 4;
    }
}
